package portscan

import (
	"errors"
	"fmt"
	"github.com/fatih/color"
	"github.com/schollz/progressbar/v3"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"
)

// ScanParams 扫描参数
type ScanParams struct {
	Ip      string      // ip或者域名
	Port    string      // 端口
	Threads int         // 协程数据
	Log     string      // 结果日志
	wait    chan string // 可以通过的IP:Port
	wg      sync.WaitGroup
	bar     *progressbar.ProgressBar
	count   chan int // 计数
	pass    []string // 所有通过的端口
}

func (s *ScanParams) initialization() ([]string, []string, int) {
	// 获取所有的ip列表
	ips, err := s.allIps()
	if err != nil {
		color.Red("解析IP列表错误: %s", err.Error())
		panic(err)
	}
	// 获取所有端口列表
	ports, err := s.allPorts()
	if err != nil {
		color.Red("解析端口列表错误: %s", err.Error())
		panic(err)
	}

	s.pass = make([]string, 0)

	// 就算每个协程需要处理的任务数据量
	count := len(ips) * len(ports)
	num := count / s.Threads

	color.Yellow("共需要扫描：%d, 启动协程数据：%d, 并发量：%d", count, s.Threads, num)

	s.wait = make(chan string, num)
	s.count = make(chan int, num/2)

	return ips, ports, count

}

func (s *ScanParams) Run() {
	s.wg = sync.WaitGroup{}

	ips, ports, count := s.initialization()
	color.Green("启动端口扫描......")

	// 测试ip:port是否可以连接
	s.wg.Add(2)
	go s.build(ips, ports)
	go s.printProgress(count)

	for i := 0; i < s.Threads; i++ {
		s.wg.Add(1)
		go s.consumer()
	}

	s.wg.Wait()
	color.Green("\n扫描结束,共有开放端口：%d个", len(s.pass))
	for _, item := range s.pass {
		color.Green("%s", item)
	}
}

func (s *ScanParams) printProgress(count int) {
	// 初始化进度值
	s.bar = progressbar.DefaultSilent(int64(count), "扫描端口:")
	for {
		if count, ok := <-s.count; ok {
			s.bar.Add(count)
			fmt.Printf("\r%s", s.bar.String())
			time.Sleep(10 * time.Millisecond)
		} else {
			break
		}
		if s.bar.IsFinished() {
			close(s.count)
			break
		}
	}
	defer s.wg.Done()
}

// allIps 获取所有的ip列表
func (s *ScanParams) allIps() ([]string, error) {
	// 判断截取是否存在类似：127.0.0.1-255这种类型的ip
	ips := strings.Split(s.Ip, "-")
	if len(ips) == 2 {
		// 说明是一个ip段，需要将ip段全部解析出来
		temp := strings.Split(ips[0], ".")
		start, _ := strconv.Atoi(temp[3])
		end, _ := strconv.Atoi(ips[1])
		if end > 255 {
			end = 255
		}
		for i := start; i < end; i++ {
			ips = append(ips, fmt.Sprintf("%s.%s.%s.%d", temp[0], temp[1], temp[2], i))
		}
	} else {
		ipAddr, err := net.ResolveIPAddr("ip", ips[0])
		if err != nil {
			return ips, errors.New(color.RedString("解析ip地址{}发生错误", ips[0]))
		}
		ips[0] = ipAddr.String()

	}
	return ips, nil
}

// allPorts 获取所以的端口
func (s *ScanParams) allPorts() ([]string, error) {
	result := make([]string, 0)
	portStr := strings.Split(s.Port, ",")
	for i := 0; i < len(portStr); i++ {
		// 判断是否存在包含的端口
		if strings.Contains(portStr[i], "-") {
			rangeStr := strings.Split(portStr[i], "-")
			start, _ := strconv.Atoi(rangeStr[0])
			end, _ := strconv.Atoi(rangeStr[1])
			for i := start; i < end; i++ {
				result = append(result, fmt.Sprintf("%d", i))
			}
		} else {
			result = append(result, portStr[i])
		}
	}
	return result, nil
}

// task 分配任务执行
// 000 000 000 000 000
func (s *ScanParams) build(ips, ports []string) {
	for i := 0; i < len(ips); i++ {
		for j := 0; j < len(ports); j++ {
			s.wait <- net.JoinHostPort(ips[i], ports[j])
		}
	}

	close(s.wait)
	defer s.wg.Done()

}

func (s *ScanParams) consumer() {
	for {
		//s.rw.Lock()
		if address, ok := <-s.wait; ok {
			conn, err := net.DialTimeout("tcp", address, time.Second*10)
			if err == nil {
				if conn != nil {
					conn.Close()
				}
				s.pass = append(s.pass, address)
			}
			s.count <- 1
		} else {
			break
		}
		//s.rw.Unlock()
	}
	defer s.wg.Done()
}
